#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Purpose: Iras network exporter
# Created: 07/06/2009
# $Id $
# Copyright (c) 2008,2009 by University College London
# Authors:
#  --<>
#
# This program is free software under the GPL (>=v2)
# Read the file COPYING coming with HydroPlatform for details.

"""
This file is part of the OpenIRAS plugin for HydroPlatform.

It allows a user to export a network to a iras.inp file
"""
from __future__ import with_statement

import codecs
import itertools
import logging
import numpy
import os


from sqlalchemy import and_, or_
from enthought.traits.api import HasTraits, Str, Enum, Int, Trait, Set

from hydroplatform.model.export import Exporter, ExportError, ExporterRegistry
from hydroplatform.model import network as nt
from hydroplatform.model import fields as fd
from hydroplatform.extension.export.common import get_value, get_objectfield

from definition import RESERVOIR_NODE, IRRIGATION_NODE, CROP_NODE, \
     HYDRO_NODE, SPILLAGE_LINK, SP_RE_LINK, RELEASE_LINK, IRRIGATION_LINK, RETURNFLOW_LINK

def get_hydro_id(node, session):
    """
    Returns the id of the hydro station linked to this node
    """
    links = session.query(nt.Link).filter(or_(nt.Link.start_node_id == node.id,
                                              nt.Link.end_node_id == node.id)).all()
    for link in links:
        # TODO : the following test could be better link.start == node, no ?
        if link.end.type.name == HYDRO_NODE and link.start.type == node.type :
            return link.end.id
        elif link.start.type.name == HYDRO_NODE and link.end.type == node.type :
            return link.start.id
    else:
        return None
        
    
def get_spill_to(node, session):
    """
    Returns the id of the node where it spills to
    """
    links = session.query(nt.Link).filter(nt.Link.start_node_id == node.id).all()
    for link in links:
        if link.type.name in [SPILLAGE_LINK, SP_RE_LINK]:
            return link.end.id
    else:
        return None

def get_release_to(node, session):
    """
    Returns the id of the node where it spills to
    """
    links = session.query(nt.Link).filter(nt.Link.start_node_id == node.id).all()
    for link in links:
        if link.type.name in [RELEASE_LINK, SP_RE_LINK]:
            return link.end.id
    else:
        return None 

def get_hydro_from(node, session):
    """
    Returns the id of the hydro node connected to this reservoir
    """
    links = session.query(nt.Link).filter(or_(nt.Link.start_node_id == node.id,
                                              nt.Link.end_node_id == node.id)).all()
    for link in links:
        if link.start.type == node.type and link.end.type.name == HYDRO_NODE :
            return link.end.id
        elif link.end.type == node.type and link.start.type.name == HYDRO_NODE :
            return link.start.id
        else:
            continue
    else:
        return None    
        
def get_water_from(node, session):
    """
    Returns the id of the reservoir linked to this node
    """
    links = session.query(nt.Link).filter(or_(nt.Link.start_node_id == node.id,
                                              nt.Link.end_node_id == node.id)).all()
    for link in links:
        if link.type.name == IRRIGATION_LINK:
            if link.start.type.name == RESERVOIR_NODE and link.end.type == node.type :
                return link.start.id
        else:
            continue
    else:
        return None
        
def return_water_to(node, session):
    """
    Returns the id of the reservoir linked to this node
    """
    links = session.query(nt.Link).filter(or_(nt.Link.start_node_id == node.id,
                                              nt.Link.end_node_id == node.id)).all()
    for link in links:
        if link.type.name == RETURNFLOW_LINK:
            if link.start.type == node.type and link.end.type.name == RESERVOIR_NODE :
                return link.end.id
        else:
            continue
    else:
        return None
        
 
    
try:
    import xlwt
    # register exporter 
    class AquaplanExporter(Exporter):
        '''
        Exports a network to the system and hydro files needed for the
        SDDP-Aquaplan model
        '''
        _name = "SDDP-Aquaplan exporter "
        _description  = '''Exports a network to Excel and csv files '''
        
        def __init__(self):
            Exporter.__init__(self)
            self.session = None
            self.network = None
        
        def export(self, network, filename, session=None):
            '''
            Export an Excel workbook to the given filename
            '''
            self.network = network
            self.basedir = os.path.dirname(filename)
            self.filename = filename
            self.session = session
            
            self.export_system_file()
            self.export_hydro_files()
            
        def export_system_file(self):
            # Create a new workbook to work on *
            self.workbook = xlwt.Workbook(encoding='cp1251')

            self.write_hpp_sheet()
            self.write_crops_sheet()
            self.write_energy_price_sheet()
            self.write_irrigation_sheet()
            self.write_reservoir_sheet()
            self.write_evaporation_sheet()
        
            # Save the file with any name *
            self.workbook.save(self.filename)
             
        def export_hydro_files(self):
            """
            Export a CSV file with the hydrology data for each hydro node
            """
            for node in self.network.nodes:
                if node.type.name == HYDRO_NODE:
                    filepath = os.path.join(self.basedir, "%s.csv" % node.id)
                    hydrodef = get_objectfield(node.type, u"hydrology", self.session)
                    hydrodata = get_value(node, u"hydrology", self.session)
                    if hydrodata is None:
                        continue
                    lst_year = None
                    year_cnt = itertools.count(1)            
                    sep = u";"
                    with open(filepath, 'w') as fh:
                        # write header 
                        headers = ["Year", "Jan", "Feb", "Mar", "Apr", 
                       "May", "Jun", "Jul", "Aug", "Sep",
                       "Oct", "Nov", "Dec"] 
                        fh.write(sep.join(headers))
                        for idx, tstamp in enumerate(hydrodef.attribute_type.dates):
                            if tstamp.year != lst_year:
                                year_idx = year_cnt.next()
                                fh.write("\n%s%s" % (tstamp.year, sep))
                                lst_year = tstamp.year
                            fh.write("%s%s" % (hydrodata[idx], sep))
        
                            
        def write_hpp_sheet(self):
            # add the hpp sheet
            hppsheet = self.workbook.add_sheet(u"Hpp")     
            # write header
            headers = ["Id", "Name", "Hydro_id", "Smin", "Smax", 
                       "Rmin", "Rmax", "rho", "Topo_l", "Topo_r",
                       "Topo_hydro", "Market_id", "Capacity [MW]",
                       "spill penalty [$/m3]", "Storage Penalty ($/m3)", "Low Storage (hm3)", "Max Head (m)"]
            for i, header in  enumerate(headers):
                hppsheet.write(0, i, header)
                
            # select the reservoir nodes
            ncnt = itertools.count(1)            
            for node in self.network.nodes:
                if node.type.name == RESERVOIR_NODE:
                    idx = ncnt.next()
                    hppsheet.write(idx, 0, node.id)
                    hppsheet.write(idx, 1, node.name)                
                    hppsheet.write(idx, 2, get_hydro_id(node, self.session))
                    hppsheet.write(idx, 3, get_value(node, u"smin", self.session) or 0)
                    hppsheet.write(idx, 4, get_value(node, u"smax", self.session) or 0)
                    hppsheet.write(idx, 5, get_value(node, u"rmin", self.session) or 0)
                    hppsheet.write(idx, 6, get_value(node, u"rmax", self.session) or 0)
                    hppsheet.write(idx, 7, get_value(node, u"rho", self.session) or 0)
                    hppsheet.write(idx, 8, get_spill_to(node, self.session))
                    hppsheet.write(idx, 9, get_release_to(node, self.session))
                    hppsheet.write(idx, 10, get_hydro_from(node, self.session))
                    hppsheet.write(idx, 11, node.id)
                    hppsheet.write(idx, 12, get_value(node, u"capacity", self.session) or 0)
                    hppsheet.write(idx, 13, get_value(node, u"spill_penalty", self.session) or 0)
                    hppsheet.write(idx, 14, get_value(node, u"storage_for_penalty", self.session) or 0)                    
                    hppsheet.write(idx, 15, get_value(node, u"penalty_for_low_storage", self.session) or 0)
                    table = get_value(node, u"reservoir", self.session)
                    if table is not None:
                        hppsheet.write(idx, 15, str(numpy.max(table[:,0])) or 0)
        
        def write_energy_price_sheet(self):
            # add the hpp sheet
            hppsheet = self.workbook.add_sheet(u"EnergyPrice")     
            # write header
            headers = ["HppId", "Jan", "Feb", "Mar", "Apr", 
                       "May", "Jun", "Jul", "Aug", "Sep",
                       "Oct", "Nov", "Dec"]
            for i, header in  enumerate(headers):
                hppsheet.write(0, i, header)
            
            ncnt = itertools.count(1)            
            for node in self.network.nodes:
                if node.type.name == RESERVOIR_NODE:
                    # get energy price
                    ts = get_value(node, u"energy_price", self.session)
                    if ts is not None:
                        nidx = ncnt.next()                        
                        hppsheet.write(nidx, 0, node.id)
                        for i in xrange(len(ts)):
                            hppsheet.write(nidx, i+1,  str(ts[i,0]))
               
        def write_crops_sheet(self):
            # add the hpp sheet
            hppsheet = self.workbook.add_sheet(u"Crops")     
            # write header
            headers = ["Id", "Name", "Yield [T/ha]", "Cropp [US $/T]",
                       "O&M costs [US $/ha]", "Return date [month", 
                       "Jan", "Feb", "Mar", "Apr", 
                       "May", "Jun", "Jul", "Aug", "Sep",
                       "Oct", "Nov", "Dec", "ky [-]", "1st irrig month"]
            for i, header in  enumerate(headers):
                hppsheet.write(0, i, header)
                
            # select the reservoir nodes
            ncnt = itertools.count(1)            
            for node in self.network.project.nodes:
                print node.name
                if node.type.name == CROP_NODE:
                    print "got a crop"
                    idx = ncnt.next()
                    hppsheet.write(idx, 0, node.id)
                    hppsheet.write(idx, 1, node.name)                
                    hppsheet.write(idx, 2, get_value(node, u"yield", self.session) or  0)
                    hppsheet.write(idx, 3, get_value(node, u"crop_price", self.session) or 0)
                    hppsheet.write(idx, 4, get_value(node, u"O&M_costs", self.session) or 0)
                    hppsheet.write(idx, 5, get_value(node, u"return_date", self.session) or 0)
                    wn = get_value(node, u"water_need", self.session)
                    if wn is not None:
                        for i in xrange(12):
                            hppsheet.write(idx, 6+i, str(wn[i,0] if not numpy.isnan(wn[i,0]) else 0))
                    hppsheet.write(idx, 18, get_value(node, u"ky", self.session))
                    hppsheet.write(idx, 19, get_value(node, u"first_irrigation_month", self.session))
        
        def write_irrigation_sheet(self):
            # add the hpp sheet
            hppsheet = self.workbook.add_sheet(u"Irrigation")     
            # write header
            headers = ["Id", "Hppid", "Cropid", "Surface [ha]", "Efficiency [-]",
                       "Return id", "Return eff [-]", "Irrig penalty [$/m3]", "Name"]
            for i, header in  enumerate(headers):
                hppsheet.write(0, i, header)
                
            # select the reservoir nodes
            ncnt = itertools.count(1)            
            for node in self.network.nodes:
                if node.type.name == IRRIGATION_NODE:
                    idx = ncnt.next()
                    hppsheet.write(idx, 0, node.id)
                    hppsheet.write(idx, 1, get_water_from(node, self.session)) 
                    hppsheet.write(idx, 2, get_value(node, u"crop", self.session))
                    hppsheet.write(idx, 3, get_value(node, u"surface", self.session))
                    hppsheet.write(idx, 4, get_value(node, u"efficiency", self.session))
                    hppsheet.write(idx, 5, return_water_to(node, self.session))
                    hppsheet.write(idx, 6, get_value(node, u"return_efficiency", self.session))
                    hppsheet.write(idx, 7, get_value(node, u"irrigation_penalty", self.session))
                    hppsheet.write(idx, 8, node.name)
                    
        def write_evaporation_sheet(self):
            # add the hpp sheet
            hppsheet = self.workbook.add_sheet(u"Evaporation")     
            # write header
            headers = ["id", "HppId", "Jan", "Feb", "Mar", "Apr", 
                       "May", "Jun", "Jul", "Aug", "Sep",
                       "Oct", "Nov", "Dec"]
            for i, header in  enumerate(headers):
                hppsheet.write(0, i, header)
            
            ncnt = itertools.count(1)            
            for node in self.network.nodes:
                if node.type.name == RESERVOIR_NODE:
                    # get energy price
                    ts = get_value(node, u"evaporation", self.session)
                    if ts is not None:
                        nidx = ncnt.next()                        
                        hppsheet.write(nidx, 0, nidx)
                        hppsheet.write(nidx, 1, node.id)
                        for i in xrange(len(ts)):
                            hppsheet.write(nidx, i+2,  str(ts[i,0]))
        
        def write_reservoir_sheet(self):
            # add the hpp sheet
            hppsheet = self.workbook.add_sheet(u"Reservoir")     
            ncnt = itertools.count(0)            
            for node in self.network.nodes:
                if node.type.name == RESERVOIR_NODE:
                    table = get_value(node, u"reservoir", self.session)
                    if table is not None:
                        nidx = ncnt.next() * 3
                        # write header
                        hppsheet.write(0, nidx    , node.id)
                        hppsheet.write(0, nidx + 1, node.id)
                        hppsheet.write(0, nidx + 2, node.id)
                        hppsheet.write(1, nidx,     u"Head [m]")
                        hppsheet.write(1, nidx + 1, u"Surface [km2]")
                        hppsheet.write(1, nidx + 2, u"Storage [hm3]") 
                        # write table
                        for cidx, row in enumerate(table):
                            hppsheet.write(cidx+1, nidx    , str(row[0]))
                            hppsheet.write(cidx+1, nidx + 1, str(row[1]))
                            hppsheet.write(cidx+1, nidx + 2, str(row[2]))
 
        
    REGISTRY = ExporterRegistry()
    REGISTRY.register(AquaplanExporter)
    logging.info("AquaplanExporter extension registered")
    
except ImportError:
    logging.warning('XLWT not available - Aquaplan export extension disabled')
    
